import ply.yacc as yacc
from lexer import tokens, lexer

# AST节点类
class ASTNode:
    def __init__(self, node_type, children=None, value=None):
        self.node_type = node_type
        self.children = children or []
        self.value = value

class Program(ASTNode):
    def __init__(self, statements):
        super().__init__('Program', statements)

class Function(ASTNode):
    def __init__(self, name, params, body, return_type=None):
        super().__init__('Function', [params, body], {'name': name, 'return_type': return_type})

class Param(ASTNode):
    def __init__(self, name, param_type):
        super().__init__('Param', [], {'name': name, 'type': param_type})

class Block(ASTNode):
    def __init__(self, statements):
        super().__init__('Block', statements)

class IfStatement(ASTNode):
    def __init__(self, condition, then_block, else_block=None):
        children = [condition, then_block]
        if else_block:
            children.append(else_block)
        super().__init__('IfStatement', children)

class WhileStatement(ASTNode):
    def __init__(self, condition, body):
        super().__init__('WhileStatement', [condition, body])

class ForStatement(ASTNode):
    def __init__(self, init, condition, update, body):
        super().__init__('ForStatement', [init, condition, update, body])

class ReturnStatement(ASTNode):
    def __init__(self, expr=None):
        super().__init__('ReturnStatement', [expr] if expr else [])

class VarDeclaration(ASTNode):
    def __init__(self, name, var_type, value=None):
        children = [value] if value else []
        super().__init__('VarDeclaration', children, {'name': name, 'type': var_type})

class Assignment(ASTNode):
    def __init__(self, target, value):
        super().__init__('Assignment', [target, value])

class BinaryOp(ASTNode):
    def __init__(self, op, left, right):
        super().__init__('BinaryOp', [left, right], {'op': op})

class UnaryOp(ASTNode):
    def __init__(self, op, expr):
        super().__init__('UnaryOp', [expr], {'op': op})

class FunctionCall(ASTNode):
    def __init__(self, name, args):
        super().__init__('FunctionCall', args, {'name': name})

class Identifier(ASTNode):
    def __init__(self, name):
        super().__init__('Identifier', [], {'name': name})

class NumberLiteral(ASTNode):
    def __init__(self, value):
        super().__init__('NumberLiteral', [], {'value': value})

class FloatLiteral(ASTNode):
    def __init__(self, value):
        super().__init__('FloatLiteral', [], {'value': value})

class BooleanLiteral(ASTNode):
    def __init__(self, value):
        super().__init__('BooleanLiteral', [], {'value': value})

# 语法分析规则 - 开始符号
def p_program(p):
    'program : statements'
    p[0] = Program(p[1])

def p_statements(p):
    '''statements : statement
                  | statements statement'''
    if len(p) == 2:
        p[0] = [p[1]]
    else:
        p[0] = p[1] + [p[2]]

def p_statement(p):
    '''statement : function_definition
                 | var_declaration
                 | assignment SEMICOLON
                 | if_statement
                 | while_statement
                 | for_statement
                 | return_statement SEMICOLON
                 | function_call SEMICOLON
                 | block'''
    p[0] = p[1]

# 函数定义
def p_function_definition(p):
    '''function_definition : DEF ID LPAREN params RPAREN COLON type block
                          | DEF ID LPAREN params RPAREN block
                          | DEF_CN ID LPAREN params RPAREN COLON type block
                          | DEF_CN ID LPAREN params RPAREN block'''
    # 确保索引安全访问
    if len(p) == 9:
        # 有返回类型的情况: DEF/DEF_CN ID LPAREN params RPAREN COLON type block
        # p[1]=DEF/DEF_CN, p[2]=ID, p[3]=LPAREN, p[4]=params, p[5]=RPAREN, p[6]=COLON, p[7]=type, p[8]=block
        p[0] = Function(p[2], p[4], p[8], p[7])
    elif len(p) == 7:
        # 没有返回类型的情况: DEF/DEF_CN ID LPAREN params RPAREN block
        # p[1]=DEF/DEF_CN, p[2]=ID, p[3]=LPAREN, p[4]=params, p[5]=RPAREN, p[6]=block
        p[0] = Function(p[2], p[4], p[6])
    else:
        raise ValueError(f"函数定义语法错误，参数长度: {len(p)}")

def p_params(p):
    '''params : param
              | params COMMA param
              | empty'''
    if len(p) == 2 and p[1] is not None:
        p[0] = [p[1]]
    elif len(p) == 4:
        p[0] = p[1] + [p[3]]
    else:
        p[0] = []

def p_param(p):
    'param : ID COLON type'
    p[0] = Param(p[1], p[3])

# 类型
def p_type(p):
    '''type : INT
            | INT_CN
            | FLOAT
            | FLOAT_CN
            | BOOL
            | BOOL_CN'''
    # 将所有类型转换为对应的小写英文类型字符串，方便后续处理
    type_map = {
        'int': 'int',
        '整数': 'int',
        'float': 'float',
        '浮点数': 'float',
        'bool': 'bool',
        '布尔': 'bool'
    }
    p[0] = type_map[p[1]]

# 变量声明
def p_var_declaration(p):
    '''var_declaration : VAR ID COLON type ASSIGN expression SEMICOLON
                      | VAR ID COLON type SEMICOLON
                      | VAR_CN ID COLON type ASSIGN expression SEMICOLON
                      | VAR_CN ID COLON type SEMICOLON'''
    if len(p) == 8:
        p[0] = VarDeclaration(p[2], p[4], p[6])
    else:
        p[0] = VarDeclaration(p[2], p[4])

# 赋值语句
def p_assignment(p):
    'assignment : ID ASSIGN expression'
    p[0] = Assignment(Identifier(p[1]), p[3])

# if语句
def p_if_statement(p):
    '''if_statement : IF LPAREN expression RPAREN block
                    | IF LPAREN expression RPAREN block ELSE block
                    | IF LPAREN expression RPAREN block ELSE if_statement
                    | IF_CN LPAREN expression RPAREN block
                    | IF_CN LPAREN expression RPAREN block ELSE block
                    | IF_CN LPAREN expression RPAREN block ELSE if_statement
                    | IF LPAREN expression RPAREN block ELSE_CN block
                    | IF LPAREN expression RPAREN block ELSE_CN if_statement
                    | IF_CN LPAREN expression RPAREN block ELSE_CN block
                    | IF_CN LPAREN expression RPAREN block ELSE_CN if_statement'''
    if len(p) == 6:
        # 简单的if语句
        p[0] = IfStatement(p[3], p[5])
    elif len(p) == 8:
        # if-else语句
        p[0] = IfStatement(p[3], p[5], p[7])
    elif len(p) == 7:
        # if-else if语句（嵌套的if_statement）
        # 将else if转换为else块中的if语句
        p[0] = IfStatement(p[3], p[5], p[6])

# while语句（支持英文和中文关键词）
def p_while_statement(p):
    '''while_statement : WHILE LPAREN expression RPAREN block
                      | WHILE_CN LPAREN expression RPAREN block'''
    p[0] = WhileStatement(p[3], p[5])

# for语句
def p_for_statement(p):
    '''for_statement : FOR LPAREN assignment SEMICOLON expression SEMICOLON assignment RPAREN block
                    | FOR_CN LPAREN assignment SEMICOLON expression SEMICOLON assignment RPAREN block'''
    p[0] = ForStatement(p[3], p[5], p[7], p[9])

# return语句
def p_return_statement(p):
    '''return_statement : RETURN expression
                        | RETURN
                        | RETURN_CN expression
                        | RETURN_CN'''
    if len(p) == 3:
        p[0] = ReturnStatement(p[2])
    else:
        p[0] = ReturnStatement()

# 代码块
def p_block(p):
    '''block : LBRACE statements RBRACE
             | LBRACE RBRACE'''
    if len(p) == 4:
        p[0] = Block(p[2])
    else:
        # 空语句块，创建一个空的Block节点
        p[0] = Block([])

# 函数调用
def p_function_call(p):
    'function_call : ID LPAREN arguments RPAREN'
    p[0] = FunctionCall(p[1], p[3])

def p_arguments(p):
    '''arguments : expression
                 | arguments COMMA expression
                 | empty'''
    if len(p) == 2 and p[1] is not None:
        p[0] = [p[1]]
    elif len(p) == 4:
        p[0] = p[1] + [p[3]]
    else:
        p[0] = []

# 表达式 - 布尔运算
def p_expression_or(p):
    'expression : expression OR expression'
    p[0] = BinaryOp('OR', p[1], p[3])

def p_expression_and(p):
    'expression : expression AND expression'
    p[0] = BinaryOp('AND', p[1], p[3])

def p_expression_comparison(p):
    '''expression : expression EQUALS expression
                 | expression NOT_EQUALS expression
                 | expression LESS expression
                 | expression LESS_EQUALS expression
                 | expression GREATER expression
                 | expression GREATER_EQUALS expression'''
    p[0] = BinaryOp(p[2], p[1], p[3])

# 表达式 - 加减
def p_expression_additive(p):
    '''expression : expression PLUS expression
                 | expression MINUS expression'''
    p[0] = BinaryOp(p[2], p[1], p[3])

# 表达式 - 乘除模
def p_expression_multiplicative(p):
    '''expression : expression TIMES expression
                 | expression DIVIDE expression
                 | expression MODULO expression'''
    p[0] = BinaryOp(p[2], p[1], p[3])

# 表达式 - 一元运算
def p_expression_unary(p):
    '''expression : NOT expression
                 | MINUS expression'''
    p[0] = UnaryOp(p[1], p[2])

# 表达式 - 括号
def p_expression_group(p):
    'expression : LPAREN expression RPAREN'
    p[0] = p[2]

# 表达式 - 原子
def p_expression_identifier(p):
    'expression : ID'
    p[0] = Identifier(p[1])

def p_expression_number(p):
    'expression : NUMBER'
    p[0] = NumberLiteral(p[1])

def p_expression_float(p):
    'expression : FLOAT_NUMBER'
    p[0] = FloatLiteral(p[1])

def p_expression_boolean(p):
    '''expression : TRUE
                  | TRUE_CN
                  | FALSE
                  | FALSE_CN'''
    p[0] = BooleanLiteral(p[1] in ('true', 'TRUE_CN'))

def p_expression_function_call(p):
    'expression : function_call'
    p[0] = p[1]

# 空产生式
def p_empty(p):
    'empty : '
    p[0] = None

# 错误处理
def p_error(p):
    if p:
        print(f"语法错误在第 {p.lineno} 行, 标记 '{p.value}', 类型 '{p.type}'")
        print(f"当前解析栈: {parser.symstack}")
        # 尝试恢复解析，跳过当前标记并继续
        # 获取当前标记的位置
        token = parser.token()  # 获取下一个标记
        if token:
            print(f"跳过标记: '{token.value}' 类型 '{token.type}' 继续解析")
            parser.errok()  # 告诉解析器我们已经处理了错误
            return token  # 返回下一个标记，继续解析
    else:
        print("语法错误: 意外的文件结束")
        print(f"当前解析栈: {parser.symstack}")

# 构建语法分析器
parser = yacc.yacc()

def parse(code):
    """对代码进行语法分析，构建AST"""
    # 重置lexer的行号，确保每次解析都从行号1开始
    lexer.lineno = 1
    return parser.parse(code, lexer=lexer)